﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using System.Diagnostics;

namespace StaticView.Algorithm
{
    class ExtendedKalmanFilter
    {
        // Parameter Count
        public int StateCount { get; private set; }
        public int ControlCount { get; private set; }

        // Current value
        public float[] StateCurrentValue { get; private set; }
        public Matrix StateCurrentNoise { get; private set; }

        // User defined covariance
        public Matrix ControlNoise { get; private set; }
        List<Matrix> ObservationNoises { get; set; }

        // Functions
        StateTransition Transition;
        List<ObservationModel> Models;

        // Essential variables
        long lastTime;

        // x, u Derivative, use in Jacobian
        float[] dx;
        float[] du;

        public ExtendedKalmanFilter(int statecount, int controlcount,
            float[] initstatevalue, Matrix initstatenoise,
            StateTransition transition, Matrix controlnoise)
        {
            #region Assignation
            StateCount = statecount;
            ControlCount = controlcount;

            StateCurrentValue = initstatevalue;
            StateCurrentNoise = initstatenoise;

            ControlNoise = controlnoise;

            Transition = transition;

            dx = new float[StateCount];
            for (int i = 0; i < StateCount; i++)
                dx[i] = 0.001f;

            du = new float[ControlCount];
            for (int i = 0; i < ControlCount; i++)
                du[i] = 0.001f;

            ObservationNoises = new List<Matrix>();
            Models = new List<ObservationModel>();
            #endregion

            #region Error-check
            if (StateCurrentValue.Length != StateCount)
                Error("Initialize error, StateCurrentValue has invalid size");
            if (StateCurrentNoise.Row != StateCount || StateCurrentNoise.Col != StateCount)
                Error("Initialize error, StateCurrentNoise has invalid size");
            if (ControlNoise.Row != ControlCount || ControlNoise.Col != ControlCount)
                Error("Initialize error, ControlNoise has invalid size");
            #endregion

            lastTime = DateTime.Now.Ticks * 100;
        }

        // Add an observation model to the filter, return the ID of the model
        public int AddObservationModel(ObservationModel model, Matrix noise)
        {
            Models.Add(model);
            ObservationNoises.Add(noise);
            return Models.Count - 1;
        }

        // Predict process
        public void Predict(float[] control, long time)
        {
            if (control.Length != ControlCount)
                Error("Predict process error, Invalid size");

            // Update the time
            long elaspedTime = time - lastTime;
            lastTime = time;

            // Predict value
            float[] newvalue = Transition(StateCurrentValue, control, elaspedTime);

            // Predict noise
            Matrix F = JacobianF(Transition, StateCurrentValue, control, elaspedTime);
            Matrix G = JacobianG(Transition, StateCurrentValue, control, elaspedTime);
            StateCurrentNoise = F * StateCurrentNoise * F.Transpose() + G * ControlNoise * G.Transpose();

            // Update state values
            StateCurrentValue = newvalue;
        }

        // Update process
        public void Update(int id, float[] observation)
        {
            ObservationModel model = Models[id];

            // Calculate the innovation value
            float[] innovation = model(StateCurrentValue);
            for (int i = 0; i < innovation.Length; i++)
            {
                innovation[i] = observation[i] - innovation[i];
            }

            // The innovation covariance
            Matrix H = JacobianH(model, StateCurrentValue);
            Matrix S = H * StateCurrentNoise * H.Transpose() + ObservationNoises[id];

            // Kalman gain
            Matrix K = StateCurrentNoise * H.Transpose() * S.Invert();
            
            // Update state
            StateCurrentValue = ((new Matrix(StateCurrentValue)) + K * (new Matrix(innovation))).ToFloatArray();

            // Update the covariance
            StateCurrentNoise = (Matrix.Identity(StateCount) - K * H) * StateCurrentNoise;
        }

        // Set derivative methods
        public void setXDerivative(float[] dx)
        {
            if (dx.Length != StateCount)
                Error("Set X Derivative error, Array size is invalid");
            else
                this.dx = dx;
        }
        public void setUDerivative(float[] du)
        {
            if (du.Length != ControlCount)
                Error("Set U Derivative error, Array size is invalid");
            else
                this.du = du;
        }

        void Error(string msg)
        {
            throw new ExtendedKalmanFilterException("EKF Exception: " + msg);
        }

        // Jacobian methods
        Matrix JacobianF(StateTransition f, float[] x, float[] u, long t)
        {
            Matrix m = new Matrix(StateCount, StateCount);
            float[] f1 = f(x, u, t);

            for (int i = 0; i < StateCount; i++)
                for (int j = 0; j < StateCount; j++)
                {
                    float[] x2 = (float[])x.Clone();
                    x2[j] = x[j] + dx[j];
                    float f2 = f(x2, u, t)[i];
                    m[i, j] = (float)((f2 - f1[i]) / dx[j]);
                }

            return m;
        }
        Matrix JacobianG(StateTransition f, float[] x, float[] u, long t)
        {
            Matrix m = new Matrix(StateCount, ControlCount);
            float[] f1 = f(x, u, t);

            for (int i = 0; i < StateCount; i++)
                for (int j = 0; j < ControlCount; j++)
                {
                    float[] u2 = (float[])u.Clone();
                    u2[j] = u[j] + du[j];
                    float f2 = f(x, u2, t)[i];
                    m[i, j] = (float)((f2 - f1[i]) / du[j]);
                }

            return m;
        }
        Matrix JacobianH(ObservationModel f, float[] x)
        {
            Matrix m = new Matrix(StateCount, StateCount);
            float[] f1 = f(x);

            for (int i = 0; i < StateCount; i++)
                for (int j = 0; j < StateCount; j++)
                {
                    float[] xdx = (float[])x.Clone();
                    xdx[j] = x[j] + dx[j];
                    double f2 = f(xdx)[i];
                    m[i, j] = (float)((f2 - f1[i]) / dx[j]);
                }

            return m;
        }
    }

    class ExtendedKalmanFilterException : Exception
    {
        public ExtendedKalmanFilterException() : base() { }
        public ExtendedKalmanFilterException(string message) : base(message) { }
    }

    delegate float[] StateTransition(float[] state, float[] control, long nanos);
    delegate float[] ObservationModel(float[] state);
}
